import os
import fitz  # PyMuPDF
import json

def get_pdf_page_counts(root_dir):
    result = {}
    
    for dirpath, _, filenames in os.walk(root_dir):
        # 获取当前目录的相对路径（相对于root_dir）
        relative_path = os.path.relpath(dirpath, root_dir)
        
        # 初始化当前层级的字典
        current_dict = result
        if relative_path != '.':
            parts = relative_path.split(os.sep)
            for part in parts:
                if part not in current_dict:
                    current_dict[part] = {}
                current_dict = current_dict[part]
        
        # 处理当前目录下的PDF文件
        for filename in filenames:
            if filename.lower().endswith('.pdf'):
                pdf_path = os.path.join(dirpath, filename)
                try:
                    with fitz.open(pdf_path) as doc:
                        current_dict[filename] = doc.page_count
                except Exception as e:
                    print(f"⚠️ 错误: 无法读取 {pdf_path} - {str(e)}")
                    current_dict[filename] = "ERROR"
    
    return result

if __name__ == "__main__":
    root_directory = "."  # 当前目录（可修改为其他路径）
    output_json = "pdf_page_counts.json"  # 输出JSON文件名
    
    # 获取PDF页码信息
    pdf_info = get_pdf_page_counts(root_directory)
    
    # 写入JSON文件（确保中文路径和内容正常显示）
    with open(output_json, "w", encoding="utf-8") as f:
        json.dump(pdf_info, f, indent=4, ensure_ascii=False)
    
    print(f"✅ 结果已保存到 {output_json}")